{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up the Working Directory\n",
    "This cell is to ensure we change the directory to anomalib source code to have access to the datasets and config files. We assume that you already went through `001_getting_started.ipynb` and install the required packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from pathlib import Path\n",
    "\n",
    "from git.repo import Repo\n",
    "\n",
    "current_directory = Path.cwd()\n",
    "if current_directory.name == \"100_datamodules\":\n",
    "    # On the assumption that, the notebook is located in\n",
    "    #   ~/anomalib/notebooks/100_datamodules/\n",
    "    root_directory = current_directory.parent.parent\n",
    "elif current_directory.name == \"anomalib\":\n",
    "    # This means that the notebook is run from the main anomalib directory.\n",
    "    root_directory = current_directory\n",
    "else:\n",
    "    # Otherwise, we'll need to clone the anomalib repo to the `current_directory`\n",
    "    repo = Repo.clone_from(url=\"https://github.com/openvinotoolkit/anomalib.git\", to_path=current_directory)\n",
    "    root_directory = current_directory / \"anomalib\"\n",
    "\n",
    "os.chdir(root_directory)\n",
    "btech_dataset_root = root_directory / \"datasets/BTech\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use BeanTech Dataset via API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pylint: disable=wrong-import-position, wrong-import-order\n",
    "# flake8: noqa\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "from torchvision.transforms import ToPILImage\n",
    "\n",
    "from anomalib.data.btech import BTech, BTechDataset\n",
    "from anomalib.data.utils import InputNormalizationMethod, get_transforms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Torch Dataset"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### DataModule\n",
    "\n",
    "Anomalib data modules are based on PyTorch Lightning (PL)'s `LightningDataModule` class. This class handles all the boilerplate code related to subset splitting, and creating the dataset and dataloader instances. A datamodule instance can be directly passed to a PL Trainer which is responsible for carrying out Anomalib's training/testing/inference pipelines. \n",
    "\n",
    "In the current example, we will show how an Anomalib data module can be created for the BeanTech Dataset, and how we can obtain training and testing dataloaders from it.\n",
    "\n",
    "To create a datamodule, we simply pass the path to the root folder of the dataset on the file system, together with some basic parameters related to pre-processing and image loading:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "btech_datamodule = BTech(\n",
    "    root=btech_dataset_root,\n",
    "    category=\"01\",\n",
    "    image_size=256,\n",
    "    train_batch_size=32,\n",
    "    eval_batch_size=32,\n",
    "    num_workers=8,\n",
    "    task=\"segmentation\",\n",
    "    normalization=InputNormalizationMethod.NONE,  # don't apply normalization, as we want to visualize the images\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the illustrative purposes of the current example, we need to manually call the `prepare_data` and `setup` methods. Normally it is not necessary to call these methods explicitly, as the PL Trainer would call these automatically under the hood.\n",
    "\n",
    "`prepare_data` checks if the dataset files can be found at the specified file system location. If not, it will download the dataset and place it in the folder.\n",
    "\n",
    "`setup` applies the subset splitting and prepares the PyTorch dataset objects for each of the train/val/test subsets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "btech_datamodule.prepare_data()\n",
    "btech_datamodule.setup()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After the datamodule has been set up, we can use it to obtain the dataloaders of the different subsets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Train images\n",
    "i, data = next(enumerate(btech_datamodule.train_dataloader()))\n",
    "print(data.keys(), data[\"image\"].shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test images\n",
    "i, data = next(enumerate(btech_datamodule.test_dataloader()))\n",
    "print(data.keys(), data[\"image\"].shape, data[\"mask\"].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As can be seen above, creating the dataloaders is pretty straghtforward, which could be directly used for training/testing/inference. We could visualize samples from the dataloaders as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img = ToPILImage()(data[\"image\"][0].clone())\n",
    "msk = ToPILImage()(data[\"mask\"][0]).convert(\"RGB\")\n",
    "\n",
    "Image.fromarray(np.hstack((np.array(img), np.array(msk))))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Torch Dataset\n",
    "\n",
    "In some cases it might be desirable to create a standalone PyTorch dataset without a PL data module. For example, this could be useful for training a PyTorch model outside Anomalib, so without the use of a PL Trainer instance. In such cases, the PyTorch Dataset instance can be instantiated directly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BTechDataset??"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To create `BTechDataset` we need to create the albumentations object that applies transforms to the input image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "get_transforms??"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_size = (256, 256)\n",
    "transform = get_transforms(image_size=256, normalization=InputNormalizationMethod.NONE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Classification Task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# BTechDataset Classification Train Set\n",
    "btech_dataset_classification_train = BTechDataset(\n",
    "    root=btech_dataset_root,\n",
    "    category=\"01\",\n",
    "    transform=transform,\n",
    "    split=\"train\",\n",
    "    task=\"classification\",\n",
    ")\n",
    "btech_dataset_classification_train.setup()  # like the datamodule, the dataset needs to be set up before use\n",
    "btech_dataset_classification_train.samples.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample = btech_dataset_classification_train[0]\n",
    "print(sample.keys(), sample[\"image\"].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As can be seen above, when we choose `classification` task and `train` split, the dataset only returns `image`. This is mainly because training only requires normal images and no labels. Now let's try `test` split for the `classification` task"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# BTech Classification Test Set\n",
    "btech_dataset_classification_test = BTechDataset(\n",
    "    root=btech_dataset_root,\n",
    "    category=\"01\",\n",
    "    transform=transform,\n",
    "    split=\"test\",\n",
    "    task=\"classification\",\n",
    ")\n",
    "btech_dataset_classification_test.setup()  # like the datamodule, the dataset needs to be set up before use\n",
    "sample = btech_dataset_classification_test[0]\n",
    "print(sample.keys(), sample[\"image\"].shape, sample[\"image_path\"], sample[\"label\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "where a classification test sample returns `image`, `image_path` and `label`. `image_path` is used to extract the filename when saving images.\n",
    "\n",
    "#### Segmentation Task\n",
    "\n",
    "It is also possible to configure the BTech dataset for the segmentation task, where the dataset object returns image and ground-truth mask."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# BTech Segmentation Train Set\n",
    "btech_dataset_segmentation_train = BTechDataset(\n",
    "    root=btech_dataset_root,\n",
    "    category=\"01\",\n",
    "    transform=transform,\n",
    "    split=\"train\",\n",
    "    task=\"segmentation\",\n",
    ")\n",
    "btech_dataset_segmentation_train.setup()\n",
    "btech_dataset_segmentation_train.samples.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above dataframe stores all the necessary information regarding the dataset. `__getitem__` method returns the corresponding information depending on the task type or train/test split."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# BTech Segmentation Test Set\n",
    "btech_dataset_segmentation_test = BTechDataset(\n",
    "    root=btech_dataset_root,\n",
    "    category=\"01\",\n",
    "    transform=transform,\n",
    "    split=\"test\",\n",
    "    task=\"segmentation\",\n",
    ")\n",
    "btech_dataset_segmentation_test.setup()\n",
    "sample = btech_dataset_segmentation_test[20]\n",
    "print(sample.keys(), sample[\"image\"].shape, sample[\"mask\"].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's visualize the image and the mask..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# img = Image.open(sample[\"image_path\"]).resize(image_size)\n",
    "img = ToPILImage()(sample[\"image\"].clone())\n",
    "msk = ToPILImage()(sample[\"mask\"]).convert(\"RGB\")\n",
    "\n",
    "Image.fromarray(np.hstack((np.array(img), np.array(msk))))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "anomalib",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13 (default, Nov  6 2022, 23:15:27) \n[GCC 9.3.0]"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "ae223df28f60859a2f400fae8b3a1034248e0a469f5599fd9a89c32908ed7a84"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
